import timm
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import numpy as np
class MixNet(nn.Module):
    def __init__(self):
        super(MixNet, self).__init__()
        self.m = timm.create_model('tf_mixnet_l', pretrained=False, features_only=True,out_indices=(2,3,4))

    def forward(self,x):
        #新增y变量,用于可视化倒残差模块的处理结果
        #y = x

        x = self.m(x)
        output = x
        #新增可视化 倒残差模块的处理结果
        # print("新增可视化")
        # y = self.m.conv_stem(y)
        # y = self.m.bn1(y)
        # y = self.m.act1(y)
        # y = self.m.blocks[0](y)
        # #这个y代表进入倒残差模块之前的
        # im = np.squeeze(y.cpu().detach().numpy())
        # # [C, H, W] -> [H, W, C]
        # im = np.transpose(im, [1, 2, 0])
        #
        # # show top 12 feature maps
        # plt.figure()
        # for i in range(12):
        #     ax = plt.subplot(3, 4, i + 1)
        #     # [H, W, C]
        #     plt.imshow(im[:, :, i], cmap='gray')
        # plt.savefig('before.png')
        #
        # y = self.m.blocks[1][0](y)
        # # 这个y代表进入倒残差模块之后的
        # im = np.squeeze(y.cpu().detach().numpy())
        # # [C, H, W] -> [H, W, C]
        # im = np.transpose(im, [1, 2, 0])
        #
        # # show top 12 feature maps
        # plt.figure()
        # for i in range(12):
        #     ax = plt.subplot(3, 4, i + 1)
        #     # [H, W, C]
        #     plt.imshow(im[:, :, i], cmap='gray')
        # plt.savefig('after.png')
        #
        # print("新增可视化结束")
        # 新增可视化 倒残差模块的处理结果--结束

        # for t in range(3):
        #
        #     print(output[t].shape)
        #     im = np.squeeze(output[t].cpu().detach().numpy())
        #     # [C, H, W] -> [H, W, C]
        #     im = np.transpose(im, [1, 2, 0])
        #
        #     # show top 12 feature maps
        #     plt.figure()
        #     for i in range(12):
        #         ax = plt.subplot(3, 4, i + 1)
        #         # [H, W, C]
        #         plt.imshow(im[:, :, i])
        #     plt.savefig('特征图第'+str(t+3)+'层.png')



        return tuple(output)

if __name__ == '__main__':
    net = MixNet()
    x = torch.randn(2, 3, 224, 224)
    o = net(x)
    for i in o:
        print(i.shape)